import os
import requests
import time
import json
from PIL import Image
from io import BytesIO
import threading
from datetime import datetime

# 基础URL和API密钥设置
base_url = 'https://api-inference.modelscope.cn/'
api_key = os.environ.get("MODELSCOPE_SDK_TOKEN")
if not api_key:
    raise ValueError("未找到环境变量 MODELSCOPE_SDK_TOKEN，请先设置该变量")

# 设置保存路径
bozo = os.environ.get("BOZO") + "/Downloads/downloads/"
os.chdir(bozo)

# 通用请求头
common_headers = {
    "Authorization": f"Bearer {api_key}",
    "Content-Type": "application/json",
}

#基础提示词
basetsc="feifei,sexy girl, a photo-realistic portrait of a young asian woman,a young girl,"

# 获取提示词输入
tsc = input("请输入英文提示词：\r\n")

# 将基础提示词和用户输入拼接
final_prompt = basetsc + " " + tsc  # 用空格连接基础提示词和用户输入
# print(f"\n完整提示词: {final_prompt}")  # 打印完整提示词供用户确认


# 定义尺寸映射字典 1024
# size_mapping = {
#     1: "768x1365",   # 9:16 比例（竖屏）
#     2: "1365x768",   # 16:9 比例（横屏）
#     3: "1024x1024",  # 1:1 比例（正方形）
#     4: "1152x910",   # 4:3 比例（传统显示器）
#     5: "910x1152",   # 3:4 比例（竖向4:3）
#     6: "836x1254",   # 2:3 比例（超宽屏）
#     7: "1254x836",   # 3:2 比例（35mm胶片）
#     8: "1566x670",   # 21:9 比例（超宽屏）
# }

size_mapping = {
    1: "960x1700",   # 9:16 比例（竖屏）- 960*1707=1,638,720
    2: "1700x960",   # 16:9 比例（横屏）- 1707*960=1,638,720
    3: "1280x1280",  # 1:1 比例（正方形）- 1280*1280=1,638,400
    4: "1480x1100",  # 4:3 比例（传统显示器）- 1478*1109=1,639,102
    5: "1100x1480",  # 3:4 比例（竖向4:3）- 1109*1478=1,639,102
    6: "1080x1536",  # 2:3 比例（超宽屏）- 1045*1568=1,638,560
    7: "1536x1080",  # 3:2 比例（35mm胶片）- 1568*1045=1,638,560
    8: "1980x830",   # 21:9 比例（超宽屏）- 1955*838=1,638,290
}

# 获取用户输入的图片尺寸选项
size_option = int(input("请输入图片尺寸数字： (1➡️ 9:16、2➡️ 16:9、3➡️ 1:1、4➡️ 4:3、5➡️ 3:4、6➡️ 2:3、7➡️ 3:2、8➡️ 21:9)\r\n "))
selected_size = size_mapping[size_option]


# 模型列表：
# MAILAND/loco-flux-2.5
# MAILAND/majicflus_v1
# bozoyan/feifei 最老
# bozoyan/FK2019_FEI
# bozoyan/FK2020_FEI
# bozoyan/FK_FEIFEI
# bozoyan/F_fei -
# bozoyan/F_feifei
# bozoyan/F_feifei01 -
# bozoyan/FEI_FK -
# bozoyan/MF_feifei -
# bozoyan/MEXX_feifei -
# bozoyan/LOCO_feifei -
# bozoyan/Xuer_feifei -

# sits elegantly on the edge of a bed, wearing a form-fitting flesh-colored bra that accentuates her natural breast curves and defined cleavage. Her posture is composed and relaxed, possibly arranging her clothing or preparing to rise. The bed is neatly made with soft sheets, pillows, and blankets, creating a cozy and intimate bedroom atmosphere. Soft natural light filters into the room, perhaps morning sunlight or warm evening glow, adding gentle shadow play and warmth to the scene. Her expression is calm and natural, with the overall composition conveying a lazy yet elegant domestic lifestyle ambiance.

# 获取模型输入（支持回车使用默认值）
model_input = input("回车确认默认模型，如需要请输入模型ID：\r\n")
model = model_input if model_input.strip() != "" else "Qwen/Qwen-Image"

# 进度条控制变量
progress_running = True
progress_count = 0

# 进度条显示函数
def show_progress():
    global progress_count
    while progress_running:
        print("=", end="", flush=True)
        progress_count += 1
        time.sleep(1)

# 记录开始时间
start_time = time.time()

# 启动进度条线程
# print("正在提交任务，请稍候: ", end="")
progress_thread = threading.Thread(target=show_progress)
progress_thread.start()

try:
    # 提交异步任务
    response = requests.post(
        f"{base_url}v1/images/generations",
        headers={**common_headers, "X-ModelScope-Async-Mode": "true"},
        data=json.dumps({
            "model": model,  # ModelScope Model-Id, required
            # "prompt": tsc,
            "prompt": final_prompt,  # 使用拼接后的完整提示词
            "size": selected_size
        }, ensure_ascii=False).encode('utf-8')
    )
    response.raise_for_status()
    task_id = response.json()["task_id"]
    
    # 更新进度条提示
    print(f"\n任务{task_id}已提交，请稍后: ", end="")
    
    # 轮询任务状态
    while True:
        result = requests.get(
            f"{base_url}v1/tasks/{task_id}",
            headers={**common_headers, "X-ModelScope-Task-Type": "image_generation"},
        )
        result.raise_for_status()
        data = result.json()
        
        if data["task_status"] == "SUCCEED":
            # 下载并处理图片
            image_url = data["output_images"][0]
            image = Image.open(BytesIO(requests.get(image_url).content))
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"魔搭_{timestamp}.jpg"
            print(f"\n图片URL地址: {image_url}")
            image.save(filename)
            
            # 打开图片
            # os.system('open ' + bozo + filename)
            break
            
        elif data["task_status"] == "FAILED":
            print("\n图片生成失败!")
            break
            
        time.sleep(5)  # 每5秒查询一次状态

finally:
    # 停止进度条线程
    progress_running = False
    progress_thread.join()
    
    # 计算并显示耗时
    end_time = time.time()
    elapsed_time = end_time - start_time
    print(f"\n图片生成完成! 总耗时: {elapsed_time:.2f} 秒")
    print(f"进度条显示时间: {progress_count} 秒")
